# -*- coding: utf-8 -*-
"""
横盘整理筛选器
筛选规则：
1. 股票最近30天横盘整理（价格在一定范围内波动，期间可以有一定的小幅起伏）
2. 成交量渐渐萎缩
3. 5日成交量均线和10日成交量均线都在下移
4. 当前5日成交量均线在10日成交量均线之下

这种形态特征可能预示着：
1. 市场关注度降低，多空力量平衡
2. 可能处于整理阶段，蓄势待发
3. 成交量萎缩可能预示变盘在即
4. 需要结合其他技术指标综合判断

使用方法：
    python horizontal_consolidation_screener.py              # 筛选所有股票
    python horizontal_consolidation_screener.py --limit 100  # 限制筛选数量（测试用）
    python horizontal_consolidation_screener.py --stock 000001  # 筛选指定股票
    python horizontal_consolidation_screener.py --days 30    # 自定义检查天数（默认30天）
"""
import logging
import pandas as pd
import numpy as np
from datetime import datetime
import sys
import os
import argparse

# 添加父目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database import DatabaseManager, StockInfoDAO
from config import DATABASE_CONFIG
import warnings

# 忽略pandas警告
warnings.filterwarnings('ignore', category=pd.errors.SettingWithCopyWarning)

# 配置日志
# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOG_DIR = os.path.join(PROJECT_ROOT, 'logs')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOG_DIR, 'horizontal_consolidation_screener.log'), encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


class HorizontalConsolidationScreener:
    """横盘整理筛选器"""
    
    def __init__(self):
        self.db = DatabaseManager(DATABASE_CONFIG)
        self.stock_dao = StockInfoDAO(self.db)
    
    def get_stock_codes(self, limit: int = None) -> list:
        """获取所有活跃股票代码"""
        try:
            sql = "SELECT stock_code, stock_name FROM new_stock_info WHERE status = 'active' ORDER BY stock_code"
            if limit:
                sql += f" LIMIT {limit}"
            
            results = self.db.execute_query(sql)
            return [(row['stock_code'], row['stock_name']) for row in results]
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return []
    
    def get_stock_daily_data(self, stock_code: str, days: int = 30) -> pd.DataFrame:
        """
        获取股票最近N天的日线数据
        
        Args:
            stock_code: 股票代码
            days: 获取天数（默认30天）
            
        Returns:
            包含交易数据的DataFrame
        """
        sql = """
        SELECT trade_date, open_price, high_price, low_price, close_price, 
               change_pct, volume, amount, turnover_rate
        FROM new_stock_daily 
        WHERE stock_code = %s 
        ORDER BY trade_date DESC 
        LIMIT %s
        """
        try:
            results = self.db.execute_query(sql, (stock_code, days))
            if not results:
                return pd.DataFrame()
            
            df = pd.DataFrame(results)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            # 转换数值列
            df['open_price'] = pd.to_numeric(df['open_price'], errors='coerce')
            df['high_price'] = pd.to_numeric(df['high_price'], errors='coerce')
            df['low_price'] = pd.to_numeric(df['low_price'], errors='coerce')
            df['close_price'] = pd.to_numeric(df['close_price'], errors='coerce')
            df['change_pct'] = pd.to_numeric(df['change_pct'], errors='coerce')
            df['volume'] = pd.to_numeric(df['volume'], errors='coerce')
            df['amount'] = pd.to_numeric(df['amount'], errors='coerce')
            df['turnover_rate'] = pd.to_numeric(df['turnover_rate'], errors='coerce')
            
            # 按日期升序排列
            df = df.sort_values('trade_date').reset_index(drop=True)
            
            # 计算成交量均线
            df['volume_ma5'] = df['volume'].rolling(window=5, min_periods=1).mean()
            df['volume_ma10'] = df['volume'].rolling(window=10, min_periods=1).mean()
            
            return df
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 数据失败: {e}")
            return pd.DataFrame()
    
    def check_horizontal_consolidation(self, df: pd.DataFrame, check_days: int = 30) -> dict:
        """
        检查股票是否满足横盘整理条件
        
        条件：
        1. 最近30天横盘整理（价格波动幅度不超过一定范围，允许小幅起伏）
        2. 成交量渐渐萎缩
        3. 5日成交量均线和10日成交量均线都在下移
        4. 当前5日成交量均线在10日成交量均线之下
        
        Args:
            df: 股票日线数据DataFrame（必须包含volume_ma5和volume_ma10列）
            check_days: 检查的天数（默认30天）
            
        Returns:
            符合条件返回详细信息字典，否则返回None
        """
        if df.empty or len(df) < check_days:
            return None
        
        # 检查必要的数据是否有效
        if df['close_price'].isna().any() or df['volume'].isna().any():
            return None
        
        # 需要至少10天数据来计算均线
        if len(df) < 10:
            return None
        
        # 计算价格波动范围
        price_range = df['close_price'].max() - df['close_price'].min()
        price_avg = df['close_price'].mean()
        price_volatility = (price_range / price_avg) * 100  # 价格波动百分比
        
        # 条件1: 横盘整理 - 价格波动幅度不超过15%（可以调整这个阈值）
        MAX_VOLATILITY = 15.0  # 最大波动幅度15%
        if price_volatility > MAX_VOLATILITY:
            return None
        
        # 条件2和3: 检查成交量均线下移趋势
        # 取最近10天的均线数据来分析下移趋势
        recent_data = df.tail(10).copy()
        
        # 计算均线下移趋势 - 使用线性回归斜率判断
        if len(recent_data) < 5:
            return None
        
        # 计算5日均线和10日均线的趋势（斜率）
        ma5_trend = np.polyfit(range(len(recent_data)), recent_data['volume_ma5'].values, 1)[0]
        ma10_trend = np.polyfit(range(len(recent_data)), recent_data['volume_ma10'].values, 1)[0]
        
        # 均线应该在下移（斜率为负）
        if ma5_trend >= 0 or ma10_trend >= 0:
            return None
        
        # 条件4: 当前5日成交量均线在10日成交量均线之下
        latest_ma5 = recent_data['volume_ma5'].iloc[-1]
        latest_ma10 = recent_data['volume_ma10'].iloc[-1]
        
        if latest_ma5 >= latest_ma10:
            return None
        
        # 检查成交量整体萎缩趋势
        # 计算最近10天的平均成交量与之前10天的平均成交量对比
        if len(df) >= 20:
            earlier_avg_volume = df.iloc[-20:-10]['volume'].mean()
            recent_avg_volume = df.iloc[-10:]['volume'].mean()
            volume_shrink_pct = ((recent_avg_volume - earlier_avg_volume) / earlier_avg_volume) * 100
        else:
            earlier_avg_volume = df.iloc[:len(df)//2]['volume'].mean()
            recent_avg_volume = df.iloc[len(df)//2:]['volume'].mean()
            volume_shrink_pct = ((recent_avg_volume - earlier_avg_volume) / earlier_avg_volume) * 100 if earlier_avg_volume > 0 else 0
        
        # 成交量应该萎缩（至少萎缩5%）
        if volume_shrink_pct > -5:
            return None
        
        # 获取最新数据
        latest_row = df.iloc[-1]
        
        # 计算统计信息
        highest_price = df['close_price'].max()
        lowest_price = df['close_price'].min()
        current_price = latest_row['close_price']
        price_change_from_high = ((current_price - highest_price) / highest_price) * 100
        price_change_from_low = ((current_price - lowest_price) / lowest_price) * 100
        
        # 计算最近30天的价格变化
        first_close = df.iloc[0]['close_price']
        last_close = latest_row['close_price']
        total_change = ((last_close - first_close) / first_close * 100)
        
        # 计算均线差值
        ma5_ma10_diff = ((latest_ma5 - latest_ma10) / latest_ma10) * 100
        
        return {
            'trade_date': latest_row['trade_date'],
            'close_price': round(current_price, 2),
            'highest_price': round(highest_price, 2),
            'lowest_price': round(lowest_price, 2),
            'price_volatility': round(price_volatility, 2),
            'price_change_from_high': round(price_change_from_high, 2),
            'price_change_from_low': round(price_change_from_low, 2),
            'total_change': round(total_change, 2),
            'latest_volume': round(latest_row['volume'], 2),
            'latest_volume_ma5': round(latest_ma5, 2),
            'latest_volume_ma10': round(latest_ma10, 2),
            'ma5_trend': round(ma5_trend, 2),
            'ma10_trend': round(ma10_trend, 2),
            'ma5_ma10_diff': round(ma5_ma10_diff, 2),
            'volume_shrink_pct': round(volume_shrink_pct, 2),
            'earlier_avg_volume': round(earlier_avg_volume, 2),
            'recent_avg_volume': round(recent_avg_volume, 2),
            'check_days': check_days
        }
    
    def screen_stocks(self, stock_codes: list = None, limit: int = None, check_days: int = 30) -> pd.DataFrame:
        """
        筛选符合条件的股票
        
        Args:
            stock_codes: 指定要筛选的股票列表，格式为[(code, name), ...]
            limit: 限制筛选的股票数量
            check_days: 检查的天数（默认30天）
            
        Returns:
            包含筛选结果的DataFrame
        """
        if stock_codes is None:
            stock_codes = self.get_stock_codes(limit=limit)
        
        if not stock_codes:
            logger.warning("没有找到需要筛选的股票")
            return pd.DataFrame()
        
        logger.info(f"开始筛选 {len(stock_codes)} 只股票（检查最近{check_days}天）...")
        
        results = []
        total = len(stock_codes)
        
        for idx, (stock_code, stock_name) in enumerate(stock_codes, 1):
            try:
                # 每10只股票输出一次进度
                if idx % 10 == 0 or idx == total:
                    logger.info(f"进度: {idx}/{total} ({idx/total*100:.1f}%)")
                
                # 获取日线数据
                df = self.get_stock_daily_data(stock_code, days=check_days)
                
                if df.empty or len(df) < check_days:
                    continue
                
                # 检查是否满足条件
                result = self.check_horizontal_consolidation(df, check_days=check_days)
                
                if result:
                    result['stock_code'] = stock_code
                    result['stock_name'] = stock_name
                    results.append(result)
                    
                    logger.info(
                        f"✓ {stock_code} {stock_name} - "
                        f"价格: {result['close_price']:.2f}, "
                        f"波动: {result['price_volatility']:.2f}%, "
                        f"MA5: {result['latest_volume_ma5']:.0f}, "
                        f"MA10: {result['latest_volume_ma10']:.0f}, "
                        f"缩量: {result['volume_shrink_pct']:.2f}%"
                    )
            
            except Exception as e:
                logger.error(f"处理股票 {stock_code} {stock_name} 时出错: {e}")
                continue
        
        if not results:
            logger.info("未找到符合条件的股票")
            return pd.DataFrame()
        
        # 转换为DataFrame并排序
        df_results = pd.DataFrame(results)
        
        # 按价格波动幅度排序（从小到大，波动越小越符合横盘特征）
        df_results = df_results.sort_values('price_volatility', ascending=True)
        
        # 重新排列列的顺序
        columns_order = [
            'stock_code', 'stock_name', 'trade_date', 'close_price',
            'highest_price', 'lowest_price', 'price_volatility',
            'price_change_from_high', 'price_change_from_low', 'total_change',
            'latest_volume', 'latest_volume_ma5', 'latest_volume_ma10',
            'ma5_trend', 'ma10_trend', 'ma5_ma10_diff',
            'volume_shrink_pct', 'earlier_avg_volume', 'recent_avg_volume',
            'check_days'
        ]
        df_results = df_results[columns_order]
        
        logger.info(f"\n筛选完成！共找到 {len(df_results)} 只符合条件的股票")
        
        return df_results
    
    def save_results(self, df: pd.DataFrame, output_dir: str = '.') -> tuple:
        """
        保存筛选结果到CSV和TXT文件
        
        Args:
            df: 筛选结果DataFrame
            output_dir: 输出目录
            
        Returns:
            (csv_file, txt_file) 文件路径元组
        """
        if df.empty:
            logger.warning("没有结果需要保存")
            return None, None
        
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        csv_file = os.path.join(output_dir, f'horizontal_consolidation_stocks_{timestamp}.csv')
        txt_file = os.path.join(output_dir, f'stock_codes_horizontal_consolidation_{timestamp}.txt')
        
        try:
            # 保存CSV文件
            df.to_csv(csv_file, index=False, encoding='utf-8-sig')
            logger.info(f"结果已保存到: {csv_file}")
            
            # 保存股票代码到TXT文件
            with open(txt_file, 'w', encoding='utf-8') as f:
                for _, row in df.iterrows():
                    f.write(f"{row['stock_code']}\n")
            logger.info(f"股票代码已保存到: {txt_file}")
            
            return csv_file, txt_file
            
        except Exception as e:
            logger.error(f"保存结果失败: {e}")
            return None, None
    
    def print_summary(self, df: pd.DataFrame):
        """打印筛选结果摘要"""
        if df.empty:
            return
        
        logger.info("\n" + "="*80)
        logger.info("筛选结果摘要")
        logger.info("="*80)
        logger.info(f"符合条件的股票数量: {len(df)}")
        logger.info(f"平均价格波动幅度: {df['price_volatility'].mean():.2f}%")
        logger.info(f"最小价格波动幅度: {df['price_volatility'].min():.2f}%")
        logger.info(f"最大价格波动幅度: {df['price_volatility'].max():.2f}%")
        logger.info(f"平均成交量萎缩幅度: {df['volume_shrink_pct'].mean():.2f}%")
        logger.info(f"平均{df['check_days'].iloc[0]}日涨幅: {df['total_change'].mean():.2f}%")
        
        logger.info(f"\n前10只波动幅度最小的股票:")
        logger.info("-"*80)
        for idx, row in df.head(10).iterrows():
            logger.info(
                f"{row['stock_code']} {row['stock_name']:8s} - "
                f"价格: {row['close_price']:7.2f} | "
                f"波动: {row['price_volatility']:5.2f}% | "
                f"MA5: {row['latest_volume_ma5']:10.0f} | "
                f"MA10: {row['latest_volume_ma10']:10.0f} | "
                f"缩量: {row['volume_shrink_pct']:6.2f}%"
            )
        logger.info("="*80)
        
        # 输出逗号分隔的股票代码
        stock_codes = df['stock_code'].tolist()
        stock_codes_str = ','.join(stock_codes)
        
        logger.info("\n" + "="*80)
        logger.info("所有匹配的股票代码（逗号分隔）:")
        logger.info("="*80)
        logger.info(stock_codes_str)
        logger.info("="*80 + "\n")


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description='横盘整理筛选器 - 筛选最近30天横盘整理且成交量萎缩的股票'
    )
    parser.add_argument(
        '--limit',
        type=int,
        help='限制筛选的股票数量（用于测试）'
    )
    parser.add_argument(
        '--stock',
        type=str,
        help='筛选指定股票代码'
    )
    parser.add_argument(
        '--days',
        type=int,
        default=30,
        help='检查的天数（默认：30天）'
    )
    parser.add_argument(
        '--output-dir',
        type=str,
        default='.',
        help='输出文件目录（默认：当前目录）'
    )
    
    return parser.parse_args()


def main():
    """主函数"""
    args = parse_arguments()
    
    logger.info("="*80)
    logger.info("横盘整理筛选器")
    logger.info("筛选规则:")
    logger.info(f"1. 股票最近{args.days}天横盘整理（价格波动幅度不超过15%）")
    logger.info("2. 成交量渐渐萎缩（最近10天平均成交量比之前萎缩至少5%）")
    logger.info("3. 5日成交量均线和10日成交量均线都在下移")
    logger.info("4. 当前5日成交量均线在10日成交量均线之下")
    logger.info("="*80 + "\n")
    
    screener = HorizontalConsolidationScreener()
    
    # 如果指定了股票代码，只筛选该股票
    if args.stock:
        stock_info = screener.stock_dao.get_stock_info_by_code(args.stock)
        if not stock_info:
            logger.error(f"未找到股票: {args.stock}")
            return
        
        stock_codes = [(stock_info['stock_code'], stock_info['stock_name'])]
        logger.info(f"筛选指定股票: {args.stock} {stock_info['stock_name']}\n")
    else:
        stock_codes = None
    
    # 执行筛选
    start_time = datetime.now()
    df_results = screener.screen_stocks(stock_codes=stock_codes, limit=args.limit, check_days=args.days)
    end_time = datetime.now()
    
    # 打印结果摘要
    screener.print_summary(df_results)
    
    # 保存结果
    if not df_results.empty:
        csv_file, txt_file = screener.save_results(df_results, output_dir=args.output_dir)
        
        if csv_file and txt_file:
            logger.info(f"\n筛选用时: {(end_time - start_time).total_seconds():.2f} 秒")
            logger.info(f"详细结果请查看: {csv_file}")
            logger.info(f"股票代码列表: {txt_file}")
    
    logger.info("\n筛选完成！")


if __name__ == '__main__':
    main()

